"""
Benchark a process given has argument. Get Private, Shared and RSS memory values and
mekes a graph with gnuplot.

@author Patrice Ferlet <metal3d@gmail.com>
@licence GNU GPL v3
"""
import sys, os, shlex, subprocess, glob
from time import time, sleep
PAGESIZE=os.sysconf("SC_PAGE_SIZE")/1024

class Benchmark:
    def __init__(self, options, prog, title, name):
        self.options = options
        self.threads = []
        self.benchname = os.path.basename(name)
        self.title = title
        self.dir = options.directory
        self.subproc = 0
        self.prog = prog
        prog = shlex.split(prog)
        self.memaverage = 0
        self.cpuaverage = 0
        self.memtrace = []
        self.cputrace = []

        #prepare time
        self.stime = time()
        p = subprocess.Popen(prog, stdout=open('/dev/null', 'w'), stderr=subprocess.STDOUT )
        #p = subprocess.Popen(prog, stderr=subprocess.STDOUT )
        print "%s is running with pid: %s" % (prog, p.pid)

        #prepare datas file to be graph
        fname = "/tmp/_datas_tograph_%s_.data" % p.pid
        self.statfile = open(fname,'w')
        #while pid is running
        while p.poll() is None:
            self.stat(p);
            #sleep( 1 / 1000000.0)
        print "done"
        self.totaltime = time() - self.stime

        print "Total time %s" % self.totaltime
        print "%d sub process was parsed" % self.subproc
        print "Memory average: %s" % (sum(self.memtrace)/len(self.memtrace))
        print "CPU average: %s" % (sum(self.cputrace)/len(self.cputrace))

        self.statfile.close()

        #now, graph !
        print "generating graph"
        out = self.graph(fname)
        os.remove(fname)
        print "%s generated" % out

    def stat(self,p):
        time2 = time() - self.stime

        (Private,Shared,Rss,CPU)= self.getStats(p.pid)

        #print stats in file
        self.statfile.write ("%s %s %s %s %s\n" % (time2, Private, Shared, Rss, CPU) )
    
        self.cputrace.append(float(CPU))
        self.memtrace.append(float(Rss))

    def graph(self, graphfile):

        font = os.getenv("WHICHROCKS_FONT")
        size = self.options.geometry
        if not size:
            size = "size " + os.getenv("WHICHROCKS_GEOMETRY", "640,480")
        else:
            size = "size "+size

        if font is not None:
            font = "font \""+font+"\""
        else:
            font = ""

        #Correct progname
        self.prog = os.path.dirname(os.path.realpath(self.prog)+'/').split(os.path.sep)[-1]
        
        title = "Benchmark %s" % self.prog.replace("_","\\_")
        if self.title is not None:
            title = self.title
        
        #create result dir
        if not os.path.exists(self.dir):
            os.mkdir(self.dir)

        out= self.dir + '/' + self.benchname +"-"+ self.prog.strip().replace(' ','-').replace('/','-') + ".png"
        p = """
set terminal png  nocrop enhanced %s %s
set title '%s'."\\ntotal time: %.4f seconds - %d detected subthreads\\nmem avegare: %s ko, cpu average: %s%%"
set autoscale y
set ytics nomirror
set ylabel "Memory usage in Ko"
set xlabel "Time in second (approximative)"

plot "%s" using 1:2  title "Private" with lines smooth csplines
replot "%s" using 1:3 title "Shared" with lines smooth csplines
replot "%s" using 1:4 title "RSS"    with lines smooth csplines

set y2tics
set autoscale y2
set y2label("CPU %%")

replot "%s" using 1:5 title "CPU" axes x1y2   with lines smooth csplines
set output "%s"
replot
        """
        p =  p % ( font, size, title , float(self.totaltime), int(self.subproc), round(sum(self.memtrace)/len(self.memtrace),2), round(sum(self.cputrace)/len(self.cputrace),2), graphfile, graphfile, graphfile, graphfile, out )
        gnuplot = subprocess.Popen(["gnuplot"], stdin=subprocess.PIPE, stdout=subprocess.PIPE)
        gnuplot.communicate(p)
        return out

    def getCPUStat(self, pid):
        #get pid cpu... this is a test... I don't find how to use /proc fs to get this
        pscom = ['ps', 'h', '-p', str(pid), '-o', 'pcpu'];
        pcpu = subprocess.Popen(pscom, stdout=subprocess.PIPE)
        pcpuv = pcpu.stdout.read()
        pcpuv = pcpuv.replace("\n","");
        pcpuv = pcpuv.replace(" ","");

        if pcpuv:
            pcpuv = float(pcpuv)
        else: 
            pcpuv = 0

        return pcpuv


    def getStats(self, pid):
        (p, s, rss) = self._getMemStats(pid)
        cpu = self.getCPUStat(pid)

        threads = []

        #fetch children process
        ppid = ['ps', 'h', '--ppid', str(pid), '-o', 'pid'];
        ppid = subprocess.Popen(ppid, stdout=subprocess.PIPE);
        
        parents = ppid.stdout.readlines() 
        for parent in parents:
            parent = parent.replace("\n", '');
            parent = parent.replace(' ','');
            if parent not in self.threads:
                self.threads.append(parent)
                print "Child detected: %s" % str(parent)
                self.subproc += 1
            if parent<>os.getpid() and os.path.isdir("/proc/%s" % str(parent)):
                (p2, s2, rss2, cpu2) = self.getStats(parent)
                cpu += cpu2
                p += float(p2)
                s += float(s2)
                rss += float(rss2)

        return (p, s, rss, cpu)


    def _getMemStats(self, pid):
        
        l = open("/proc/"+str(pid)+"/statm", "rt").readline().split()

        Rss=int(l[1])*PAGESIZE
        Shared=int(l[2])*PAGESIZE
        Private = Rss-Shared
        return (Private, Shared, Rss)



if __name__ == "__main__":
    prog = " ".join(sys.argv[1:])
    if len(prog) == 0:
        print "You must provide a command to bench"
        sys.exit(1)
    #do benchmark
    Benchmark('./', prog, None)
    sys.exit(0)

